from __future__ import (absolute_import, division, print_function, unicode_literals)
import sys

#sys.path.append("D:\\pycharm\\workspace\\python-demo\\test")
sys.path.append("/Users/yishanli/python/myProject/python-demo/test/")
import backtrader as bt
import pandas as pd
from tushares.TuShareTool import TuShareTool  # 导入用于获取股票数据的TuShareTool

# 初始化 TuShareTool
tushare = TuShareTool()


def get_data(code='600483.SH', start_date='2022-10-01', end_date='2023-10-01'):
    # 使用TuShareTool获取指定股票代码和日期范围内的股票数据
    df = tushare.get_stock_data(code, start_date, end_date)

    # 调整数据格式，将日期列转换为Datetime类型并设置为索引
    df['Datetime'] = pd.to_datetime(df['trade_date'])
    df.set_index("Datetime", inplace=True)

    # 添加一个名为'openinterest'的列，值都设置为0
    df['openinterest'] = 0

    # 重命名'vol'列为'volume'
    df.rename(columns={'vol': 'volume'}, inplace=True)

    # 选择并返回特定的列，同时反转数据，以便按照时间升序排列
    df = df[['open', 'high', 'low', 'close', 'volume', 'openinterest']]
    return df.iloc[::-1]  # 反转数据，按时间升序排列


class TestStrategy(bt.Strategy):
    params = (
        ('maperiod', 15),
    )

    def log(self, txt, dt=None):
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    def __init__(self):
        self.dataclose = self.datas[0].close
        self.order = None
        self.buyprice = None
        self.buycomm = None
        self.sma = bt.indicators.SimpleMovingAverage(self.datas[0], period=self.params.maperiod)  # 创建移动均线指标

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(
                    'BUY EXECUTED,Price: %.2f,Cost:%.2f,Comm %.2f' %
                    (order.executed.price,
                     order.executed.value,
                     order.executed.comm))
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            else:
                self.log('SELL EXECUTED,Price: %.2f,Cost:%.2f,Comm %.2f' %
                         (order.executed.price,
                          order.executed.value,
                          order.executed.comm))

            self.bar_executed = len(self)
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')
        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        self.log('OPERATION PROFIT,GROSS %.2f,NET %.2f' %
                 (trade.pnl, trade.pnlcomm))

    def next(self):
        self.log('Close, %.2f' % self.dataclose[0])
        if self.order:
            return
        if not self.position:
            if self.datas[0].close[0] > self.sma[0]:  # 如果收盘价上穿均线
                self.log('BUY CREATE, %.2f' % self.dataclose[0])
                self.order = self.buy(size=1000)  # 发出买入指令，买入指定数量的股票
        else:
            if self.datas[0].close[0] < self.sma[0]:  # 如果收盘价下穿均线
                self.log('SELL CREATE, %.2f' % self.dataclose[0])
                self.order = self.sell(size=1000)


if __name__ == '__main__':
    # 获取数据
    stock_df = get_data("603019.SH", '20240101', '20240801')  # 获取特定股票的数据
    # 初始化 Backtrader 大脑
    cerebro = bt.Cerebro()  # 创建Backtrader大脑

    # 添加策略
    cerebro.addstrategy(TestStrategy, maperiod=15)  # 向大脑添加自定义策略

    # 添加数据
    data = bt.feeds.PandasData(dataname=stock_df)  # 创建数据源，使用Pandas数据
    cerebro.adddata(data, "pingan")  # 添加数据到大脑，并命名为"pingan"
    start_cash = 100000  # 设置初始资金
    # 添加绘图
    cerebro.addsizer(bt.sizers.FixedSize, stake=10)

    # 添加分析
    cerebro.addanalyzer(bt.analyzers.PyFolio, _name='pyfolio')  # 添加PyFolio分析器

    cerebro.broker.setcash(start_cash)
    cerebro.broker.setcommission(commission=0.001)

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    cerebro.run()
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
    # 绘制均线和股价图
    cerebro.plot()

"""
bt.indicators.ExponentialMovingAverage(self.datas[0], period=25)
bt.indicators.WeightedMovingAverage(self.datas[0], period=25).subplot = True
bt.indicators.StochasticSlow(self.datas[0])
bt.indicators.MACDHisto(self.datas[0])
rsi = bt.indicators.RSI(self.datas[0])
bt.indicators.SmoothedMovingAverage(rsi, period=10)
bt.indicators.ATR(self.datas[0]).plot = False
"""
